from sklearn.metrics import roc_curve, roc_auc_score
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
import scipy
import shelve

class img_tool():
    def __init__(self) -> None:
        pass

    #   这里用来绘制roc曲线图
    #   感觉是需要保存预测值和实际值即可
    # ------------------------------ auc ---------------------------------------------------------
    def draw_roc(self):
        modelName=["tlsh","functionSim","img","MGMN","siamese_graphsage"]
        plt.figure(figsize = (7,7))
        for i in range(len(modelName)):
            print("生成{}模型的roc曲线".format(modelName[i]))
            with shelve.open("/home/chenyongwei/function_sim_project/all_data/indicators/auc/{}_auc".format(modelName[i])) as file:
                pred=file["pred"]
                tar=file["true"]
            fpr, tpr, thresholds = roc_curve(tar, pred)
            auc = roc_auc_score(tar, pred)
            plt.plot(fpr, tpr, label='{} auc: {:.2f}'.format(modelName[i],auc))
        plt.plot([0, 1], [0, 1], 'k--')  # 绘制对角线
        plt.xlabel('False Positive Rate')
        plt.ylabel('True Positive Rate')
        plt.title('ROC')
        plt.legend(loc='lower right')
        plt.savefig('/home/chenyongwei/function_sim_project/all_data/img/roc_curve.png', dpi=1000)
        plt.show()

    def get_auc_data(self):
        modelName=["tlsh","img","functionSim","mgmn","mgmn_without_matching"]
        res=[]
        tar=[]
        for i in modelName:
            file =shelve.open("/home/chenyongwei/project/malSimModels/dataInf/auc/{}_auc".format(i))
            res.append(np.array(file["pred"]))
            tar.append(np.array(file["actual"]))
            file.close()
        return res,tar

    #   这个先不管吧，感觉没什么用
    #   绘制模型的下降曲线,
    # ------------------------------ loss curve ---------------------------------------------------------
    def get_loss_data(self,a):
        res=[]
        with open(a) as file:
            data=file.readlines()
            for i in range(len(data)):
                end=data[i].index(",")
                res.append(float(data[i][7:end]))
        return res

    def draw_loss(self):
        paths=["/home/chenyongwei/project/malSimModels/dataInf/loss/functionSim_loss.txt",\
            "/home/chenyongwei/project/malSimModels/dataInf/loss/img_loss.txt",\
                "/home/chenyongwei/project/malSimModels/dataInf/loss/mgmn_loss.txt",\
                "/home/chenyongwei/project/malSimModels/dataInf/loss/mgmn_without_matching_loss.txt"]
        names=["functionSim","img","MGMN","MGMG_without_match"]

        plt.figure(figsize = (20,4))
        for i in range(len(paths)):
            data=self.get_loss_data(paths[i])
            x = np.array(list(range(len(data))))    
            y = np.array(data)    
            plt.subplot(1,4,i+1)
            y_smooth = scipy.signal.savgol_filter(y,30,5)
            plt.plot(x, y,alpha=0.4,linewidth =2.5,label = "{} Original curve".format(names[i]))   
            plt.plot(x, y_smooth,alpha=1,label = "{} fit curve".format(names[i]))   
            plt.legend(loc='best')
            plt.title("{} loss".format(names[i]),fontsize = 10,color = 'blue',fontweight = "heavy",loc ='left')
        plt.savefig("/home/chenyongwei/project/malSimModels/modelImg/all_loss")
    
    def draw_confusion_matrix(self,name,matrix,lables):
        # 创建一个 Matplotlib 图表
        plt.figure(figsize=(20, 10),dpi=600)
        # 创建表格并添加数据
        table = plt.table(cellText=matrix, loc='center')
        # 设置表格的样式
        table.auto_set_font_size(False)
        table.set_fontsize(6)
        plt.tight_layout()
        plt.axis('off')

        table_props = table.properties()['celld']

        # # 设置第一行的字体样式
        # for cell in table_props.values():
        #     cell.set_fontsize(2)
        #     cell.get_text().set_weight('bold')
        #     cell.set_facecolor('#C3E6E3')

        # 设置第一列的字体样式
        for i in range(1, len(matrix)):
            cell = table_props[(i, 0)]
            cell.set_fontsize(2)
            cell = table_props[(0, i)]
            cell.set_fontsize(2)

        # 自定义颜色映射
        cmap = mcolors.LinearSegmentedColormap.from_list("", ["white", "darkblue"])
        # 设置单元格颜色
        newMatrix=[row[1:] for row in matrix[1:]]
        maxvalue=np.max(newMatrix)
        for i in range(1,len(matrix)):
            for j in range(1,len(matrix[0])):
                value = matrix[i][j] + 150 if matrix[i][j]!=0 else 0  # 数据值
                value=min(value,maxvalue)
                cell = table[(i, j)]  # 获取单元格
                cell.set_facecolor(cmap(value/ maxvalue))  # 根据数据值设置颜色
                if i==j:
                    cell.set_facecolor('#C3E6C3')

        plt.savefig("/home/chenyongwei/function_sim_project/all_data/img/{}_cinfusion_matrix.png".format(name))

if __name__=="__main__":
    gene_img=img_tool()
    # gene_img.draw_roc()
    gene_img.draw_confusion_matrix(0,0,0)